package com.webcamtracker.color;

import com.webcamtracker.color.segmentation.Segment;
import com.webcamtracker.color.segmentation.SizeSegmentComparator;
import com.webcamtracker.configuration.RootPref;
import com.webcamtracker.image.Color;
import com.webcamtracker.image.color.rgb.RGBColor;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.PixelGrabber;
import java.util.*;
import java.util.prefs.Preferences;


public class SegmentationTrackingObjectFinder {
    // The input image and its dimensions.
    private int width, height;

    // A matrix for the pixel values, one for the selected labels, one which indicates
    // whether a pixel already has a label.
    private int[] pixels;
    private int[] labels;

    // Counters for pixels in each region
    private Map<Integer, Integer> count;
    private SortedSet<Segment> segments;


    //Map: region number > points of the region
    private Map<Color, SortedSet<Segment>> regions;

    private static Preferences pref = RootPref.getRootPref().node("segmentation");

    public Map<Color, SortedSet<Segment>> segment(BufferedImage image) {
        initVariables(image);
        //  int min = pref.getInt("min", 100);

        int numberOfRegion = 0;
        Segment currentSegment = null;

        java.util.List<Point> mustDo = new LinkedList<Point>();
        boolean segmentAdded = false;
        for (int h = 0; h < height * width; h++) {
            final int y = h / width;
            final int x = h - (y * width);

            // Is this pixel unlabeled?
            if (labels[h] < 0) {
                numberOfRegion++;
                final RGBColor color = new RGBColor(image.getRGB(x, y));
                currentSegment = new Segment(color, numberOfRegion);
                segmentAdded = false;
                mustDo.add(new Point(x, y));
                labels[h] = numberOfRegion; // label it as one on a new region
                count.put(numberOfRegion, 1);
            }
            assert currentSegment != null;
            // Check all the pixels on the stack. There may be more than one!
            while (mustDo.size() > 0) {
                Point thisPoint = mustDo.remove(0);

                // Check 8-neighborhood
                for (int th = -1; th <= 1; th++)
                    for (int tw = -1; tw <= 1; tw++) {
                        int rx = thisPoint.x + tw;
                        int ry = thisPoint.y + th;
                        // Skip pixels outside of the image.
                        if ((rx < 0) || (ry < 0) || (ry >= height) || (rx >= width)) continue;
                        if (labels[ry * width + rx] < 0) {
                            final int pixelA = pixels[ry * width + rx];
                            final int pixelB = pixels[thisPoint.y * width + thisPoint.x];

                            if (pixelA == pixelB) {
                                final Point point = new Point(rx, ry);
                                mustDo.add(point);
                                labels[ry * width + rx] = numberOfRegion;
                                currentSegment.addPoint(point);
                            }
                        }

                    } // ended neighbors checking
            } // ended stack scan

            if (!segmentAdded && currentSegment.getPoints().size() > 100) {
                addSegment(currentSegment.getColor(), currentSegment);
                segments.add(currentSegment);
                segmentAdded = true;
            }

        }

//        paintImage(image);

        return regions;
    }

    public SortedSet<Segment> getSegments() {
        return segments;
    }

    private void addSegment(Color color, Segment currentSegment) {
        SortedSet<Segment> segments = regions.get(color);
        if (segments == null) {
            segments = new TreeSet<Segment>(new SizeSegmentComparator());
            regions.put(color, segments);
        }
        segments.add(currentSegment);
    }

    private void initVariables(BufferedImage image) {
        this.width = image.getWidth();
        this.height = image.getHeight();
        this.segments = new TreeSet<Segment>(new SizeSegmentComparator());

        this.labels = new int[width * height];
        int i = 0;
        for (int i1 = 0; i1 < labels.length; i1++) {
            labels[i1] = -1;
        }

        PixelGrabber pg = new PixelGrabber(image, 0, 0, -1, -1, true);

        try {
            pg.grabPixels();
        }
        catch (InterruptedException e) {
            e.printStackTrace();
        }


        this.pixels = (int[]) pg.getPixels();

        this.count = new TreeMap<Integer, Integer>();
        this.regions = new HashMap<Color, SortedSet<Segment>>();
    }
}
